package com.example.plantuml.controller;

import com.example.plantuml.model.RequestTrace;
import com.example.plantuml.tools.CallChainManager;
import com.example.plantuml.tools.PlantUmlGenerator;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.nio.charset.StandardCharsets;

@RestController
@RequestMapping("/traces")
public class TraceController {

    @Autowired
    private CallChainManager callChainManager;

    @Autowired
    private PlantUmlGenerator plantUmlGenerator;

    /**
     * 获取所有跟踪记录ID
     */
    @GetMapping
    public ResponseEntity<?> getAllTraces() {
        return ResponseEntity.ok(callChainManager.getAllTraceIds());
    }

    /**
     * 获取指定跟踪ID的PlantUML时序图
     */
    @GetMapping("/plantuml/{traceId}")
    public ResponseEntity<Object> getTracePlantUml(@PathVariable String traceId) {
        RequestTrace trace = callChainManager.getCompletedTrace(traceId);
        if (trace == null) {
            return ResponseEntity.notFound().build();
        }
        try {
            String plantUml = plantUmlGenerator.generateSequenceDiagram(trace);
            byte[] imageBytes = plantUmlGenerator.generatePngImage(plantUml);
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.IMAGE_PNG);
            headers.setContentLength(imageBytes.length);

            return new ResponseEntity<>(imageBytes, headers, HttpStatus.OK);
//            return new ResponseEntity<>(plantUml, headers, HttpStatus.OK);
        } catch (Exception ex) {
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(new MediaType("text", "plain", StandardCharsets.UTF_8));
            return new ResponseEntity<>("异常", headers, HttpStatus.OK);

        }
    }
}
